from math import sqrt
import numpy as np
from data.trs_wr import *
from concurrent.futures import ProcessPoolExecutor

key_point=300

class Align_TSDM:
    def __init__(self,gkp=True):
        self.gkp=gkp
        self.share_dict={}

    def generate_key_point(self,trace):
        if len(trace)<max_point:
            return trace, np.array(range(len(trace)))
        amplitude = np.abs(trace[1:-1] - (trace[:-2] + trace[2:]) / 2)
        largest_delta = np.argsort(amplitude)[-key_point:]
        key_point = np.sort(largest_delta)
        key_point = np.hstack((0, key_point + 1, len(trace) - 1))
        return trace[key_point], key_point

    def TSDM(self,trace_in):
        trace_base=self.share_dict['trace_base']
        norm_base=self.share_dict['norm_base']
        trace_vec_base=self.share_dict['trace_vec_base']
        key_point_vec_base=self.share_dict['key_point_vec_base']
        if self.gkp:
            trace,key_point=self.generate_key_point(trace_in)
        else:
            trace=trace_in
            key_point=np.array(range(len(trace_in)))
        trace_vec = trace[1:] - trace[:-1]
        key_point_vec = key_point[1:] - key_point[:-1]
        distance = np.zeros((len(trace_vec) + 1, len(trace_vec_base) + 1))
        distance[0,1:],distance[1:,0]=-1,-1
        for m in range(len(trace_vec)):
            for n in range(len(trace_vec_base)):
                vec_mul = trace_vec[m] * trace_vec_base[n] + key_point_vec[m] * key_point_vec_base[n]
                norm = norm_base[n] * sqrt(trace_vec[m] ** 2 + key_point_vec[m] ** 2)
                cos_alpha = vec_mul / norm
                cos_alpha = cos_alpha if cos_alpha > 0 else 0
                distance[m + 1, n + 1] = cos_alpha + max(distance[m, n], distance[m + 1, n], distance[m, n + 1])
        x, y = len(trace_vec) - 1, len(trace_vec_base) - 1
        trace_new = np.zeros((len(trace_base)))
        trace_new[0], trace_new[y],trace_new[-1] = trace[0], trace[x],trace[-1]
        count = np.zeros((len(trace_base)),np.uint32)
        count[0], count[y],count[-1] = 1, 1, 1
        while True:
            max_value = max(distance[x, y + 1], distance[x, y], distance[x + 1, y])
            if max_value==distance[x, y + 1]:x=x-1
            elif max_value==distance[x + 1, y]:y=y-1
            else:
                x = x - 1
                y = y - 1
            trace_new[y] += trace[x]
            count[y] += 1
            if x == 0 and y == 0: break
        trace_new /= count
        return trace_new

    def align(self,traces,signal_update_process_bar=None):
        if self.gkp:
            trace_base, key_point_base=self.generate_key_point(traces[0])
        else:
            trace_base=traces[0]
            key_point_base=np.array(range(len(traces[0])))
        trace_vec_base = trace_base[1:] - trace_base[:-1]
        key_point_vec_base = key_point_base[1:] - key_point_base[:-1]
        norm_base = np.sqrt(trace_vec_base ** 2 + key_point_vec_base ** 2)
        self.share_dict['trace_base']=trace_base
        self.share_dict['norm_base']=norm_base
        self.share_dict['key_point_base']=key_point_base
        self.share_dict['trace_vec_base']=trace_vec_base
        self.share_dict['key_point_vec_base']=key_point_vec_base
        trace_align = [trace_base]
        with ProcessPoolExecutor() as pool:
            results=pool.map(self.TSDM,traces[1:])
            for i,result in enumerate(results):
                trace_align.append(result)
                if signal_update_process_bar is not None:
                    signal_update_process_bar.emit(100*i/(traces.shape[0]-1))
        return np.array(trace_align)